"""Temporal CNN量化模型"""
import pandas as pd
import numpy as np
from typing import Dict, Optional, Tuple
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv1D, MaxPooling1D, Flatten, Dense, Dropout
from tensorflow.keras.optimizers import Adam
from sklearn.preprocessing import MinMaxScaler
from .base_model import BaseQuantModel


class TemporalCNNModel(BaseQuantModel):
    """Temporal CNN时间序列模型（支持分类和回归）"""
    
    def __init__(self, task_type: str = "classification", sequence_length: int = 60, **kwargs):
        """
        初始化Temporal CNN模型
        
        Args:
            task_type: 任务类型，"classification"或"regression"
            sequence_length: 时间序列长度（输入窗口大小），默认60
            **kwargs: Keras模型参数
                常用参数：
                - filters: 卷积核数量，默认64
                - kernel_size: 卷积核大小，默认3
                - dropout_rate: Dropout比例，默认0.2
                - learning_rate: 学习率，默认0.001
                - batch_size: 批次大小，默认32
                - epochs: 训练轮数，默认50
        """
        super().__init__(model_name="TemporalCNN", task_type=task_type)
        self.sequence_length = sequence_length
        
        default_params = {
            'filters': 64,
            'kernel_size': 3,
            'dropout_rate': 0.2,
            'learning_rate': 0.001,
            'batch_size': 32,
            'epochs': 50
        }
        self.model_params = {**default_params, **kwargs}
        
        # 特征缩放器
        self.feature_scaler = MinMaxScaler()
        self.target_scaler = None  # 回归任务时使用
    
    def _build_model(self, input_shape: Tuple[int, int], **kwargs):
        """构建Temporal CNN模型"""
        params = {**self.model_params, **kwargs}
        
        model = Sequential()
        
        # 第一个卷积层
        model.add(Conv1D(
            filters=params['filters'],
            kernel_size=params['kernel_size'],
            activation='relu',
            input_shape=input_shape
        ))
        model.add(MaxPooling1D(pool_size=2))
        model.add(Dropout(params['dropout_rate']))
        
        # 第二个卷积层
        model.add(Conv1D(
            filters=params['filters'],
            kernel_size=params['kernel_size'],
            activation='relu'
        ))
        model.add(MaxPooling1D(pool_size=2))
        model.add(Dropout(params['dropout_rate']))
        
        # 展平
        model.add(Flatten())
        
        # 全连接层
        model.add(Dense(50, activation='relu'))
        model.add(Dropout(params['dropout_rate']))
        
        if self.task_type == "classification":
            model.add(Dense(1, activation='sigmoid'))
            model.compile(
                optimizer=Adam(learning_rate=params['learning_rate']),
                loss='binary_crossentropy',
                metrics=['accuracy']
            )
        else:
            model.add(Dense(1))
            model.compile(
                optimizer=Adam(learning_rate=params['learning_rate']),
                loss='mse',
                metrics=['mae']
            )
        
        self.model = model
    
    def _create_sequences(
        self,
        X: pd.DataFrame,
        y: Optional[pd.Series] = None
    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        """创建时间序列样本"""
        X_scaled = self.feature_scaler.transform(X)
        
        X_seq = []
        y_seq = [] if y is not None else None
        
        for i in range(self.sequence_length, len(X_scaled)):
            X_seq.append(X_scaled[i - self.sequence_length:i])
            if y is not None:
                y_seq.append(y.iloc[i])
        
        X_seq = np.array(X_seq)
        
        if y is not None:
            y_seq = np.array(y_seq)
            if self.task_type == "classification":
                y_seq = y_seq.reshape(-1, 1)
        
        return X_seq, y_seq
    
    def _train_model(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_val: Optional[pd.DataFrame] = None,
        y_val: Optional[pd.Series] = None,
        **kwargs
    ) -> Dict:
        """训练Temporal CNN模型"""
        # 拟合特征缩放器
        self.feature_scaler.fit(X_train)
        
        # 如果是回归任务，也需要缩放目标变量
        if self.task_type == "regression":
            self.target_scaler = MinMaxScaler()
            y_train_scaled = y_train.values.reshape(-1, 1)
            self.target_scaler.fit(y_train_scaled)
        
        # 创建时间序列
        X_train_seq, y_train_seq = self._create_sequences(X_train, y_train)
        
        # 构建模型
        input_shape = (X_train_seq.shape[1], X_train_seq.shape[2])
        self._build_model(input_shape, **kwargs)
        
        # 准备验证集
        validation_data = None
        if X_val is not None and y_val is not None:
            X_val_seq, y_val_seq = self._create_sequences(X_val, y_val)
            if self.task_type == "regression" and self.target_scaler is not None:
                y_val_seq_scaled = self.target_scaler.transform(y_val_seq.reshape(-1, 1))
                validation_data = (X_val_seq, y_val_seq_scaled.ravel())
            else:
                validation_data = (X_val_seq, y_val_seq)
        
        # 训练模型
        params = {**self.model_params, **kwargs}
        
        history = self.model.fit(
            X_train_seq,
            y_train_seq if self.task_type == "classification" else self.target_scaler.transform(y_train_seq.reshape(-1, 1)).ravel(),
            batch_size=params['batch_size'],
            epochs=params['epochs'],
            validation_data=validation_data,
            verbose=0
        )
        
        # 计算训练集得分
        train_loss = history.history['loss'][-1]
        train_metric = history.history[list(history.history.keys())[1]][-1]
        
        result = {
            'train_loss': train_loss,
            'train_metric': train_metric
        }
        
        # 如果有验证集
        if validation_data is not None:
            val_loss = history.history.get('val_loss', [None])[-1]
            val_metric_key = [k for k in history.history.keys() if k.startswith('val_') and k != 'val_loss'][0]
            val_metric = history.history.get(val_metric_key, [None])[-1]
            result['val_loss'] = val_loss
            result['val_metric'] = val_metric
        
        return result
    
    def _create_sequences_for_prediction(self, X: pd.DataFrame) -> np.ndarray:
        """为预测创建时间序列（不包含目标变量）"""
        X_scaled = self.feature_scaler.transform(X)
        
        if len(X_scaled) < self.sequence_length:
            # 如果数据不足，使用padding
            padding = np.tile(X_scaled[0], (self.sequence_length - len(X_scaled), 1))
            X_scaled = np.vstack([padding, X_scaled])
        
        # 只返回最后一个序列
        X_seq = X_scaled[-self.sequence_length:].reshape(1, self.sequence_length, -1)
        
        return X_seq
    
    def _predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        """预测概率或值"""
        # 为每个样本创建序列（简化版：使用滑动窗口）
        predictions = []
        
        for i in range(len(X)):
            # 获取从开始到当前位置的数据
            X_window = X.iloc[max(0, i - self.sequence_length + 1):i + 1]
            
            if len(X_window) < self.sequence_length:
                # 如果数据不足，使用padding
                padding_count = self.sequence_length - len(X_window)
                first_row = X_window.iloc[0:1]
                padding = pd.concat([first_row] * padding_count, ignore_index=True)
                X_window = pd.concat([padding, X_window], ignore_index=True)
            
            X_seq = self._create_sequences_for_prediction(X_window)
            
            pred = self.model.predict(X_seq, verbose=0)
            
            if self.task_type == "regression" and self.target_scaler is not None:
                # 反缩放
                pred = self.target_scaler.inverse_transform(pred.reshape(-1, 1)).ravel()
            
            predictions.append(pred[0])
        
        predictions = np.array(predictions)
        
        if self.task_type == "classification":
            # 对于分类，返回概率
            return predictions.reshape(-1)
        
        return predictions
    
    def save(self, filepath: str):
        """保存模型（重写以保存scaler）"""
        if not self.is_trained:
            raise ValueError("模型尚未训练，无法保存")
        
        import pickle
        
        # 保存Keras模型
        model_path = filepath.replace('.pkl', '_model.h5')
        self.model.save(model_path)
        
        # 保存其他信息
        model_data = {
            'model_name': self.model_name,
            'task_type': self.task_type,
            'feature_columns': self.feature_columns,
            'model_params': self.model_params,
            'is_trained': self.is_trained,
            'sequence_length': self.sequence_length,
            'feature_scaler': self.feature_scaler,
            'target_scaler': self.target_scaler,
            'model_path': model_path
        }
        
        import os
        os.makedirs(os.path.dirname(filepath) if os.path.dirname(filepath) else '.', exist_ok=True)
        
        with open(filepath, 'wb') as f:
            pickle.dump(model_data, f)
    
    def load(self, filepath: str):
        """加载模型（重写以加载scaler）"""
        import pickle
        
        with open(filepath, 'rb') as f:
            model_data = pickle.load(f)
        
        # 加载Keras模型
        self.model = keras.models.load_model(model_data['model_path'])
        
        self.model_name = model_data['model_name']
        self.task_type = model_data['task_type']
        self.feature_columns = model_data['feature_columns']
        self.model_params = model_data.get('model_params', {})
        self.is_trained = model_data['is_trained']
        self.sequence_length = model_data.get('sequence_length', 60)
        self.feature_scaler = model_data.get('feature_scaler')
        self.target_scaler = model_data.get('target_scaler')

